"""
python fujian2_transformer_predict.py
"""

import os
import json
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from torch import nn, optim
from torch.utils.data import DataLoader, Dataset
from datetime import timedelta

# 自定义时间序列数据集
class TimeSeriesDataset(Dataset):
    def __init__(self, data, sequence_length=10):
        self.data = data
        self.sequence_length = sequence_length

    def __len__(self):
        return len(self.data) - self.sequence_length

    def __getitem__(self, index):
        return (self.data[index:index + self.sequence_length],
                self.data[index + self.sequence_length])

# Transformer 模型
class TransformerModel(nn.Module):
    def __init__(self, input_size=1, hidden_dim=64, nhead=4, num_layers=2):
        super(TransformerModel, self).__init__()
        self.encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=nhead)
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        self.fc = nn.Linear(hidden_dim, 1)
        self.embedding = nn.Linear(input_size, hidden_dim)

    def forward(self, src):
        src = self.embedding(src)
        output = self.transformer_encoder(src)
        output = self.fc(output)
        return output

# 单个分类的时间序列预测
def predict_single_category(input_json_path, output_json_path, graph_output_path):
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
    os.makedirs(os.path.dirname(graph_output_path), exist_ok=True)

    # 读取JSON文件并转换为DataFrame
    with open(input_json_path, 'r') as file:
        data = json.load(file)
    df = pd.DataFrame(data)
    df['ds'] = pd.to_datetime(df['date'])
    df['y'] = df['sales']

    # 数据预处理
    values = df['y'].values
    scaler = (values.min(), values.max())
    scaled_data = (values - scaler[0]) / (scaler[1] - scaler[0])
    
    # 创建数据集和加载器
    sequence_length = 10
    dataset = TimeSeriesDataset(torch.FloatTensor(scaled_data).unsqueeze(-1), sequence_length)
    data_loader = DataLoader(dataset, batch_size=32, shuffle=True)

    # 定义模型和训练参数
    model = TransformerModel()
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    model.train()

    # 训练模型
    epochs = 20
    for epoch in range(epochs):
        for X, y in data_loader:
            X = X.permute(1, 0, 2)  # 变换维度为(seq_len, batch_size, input_size)
            optimizer.zero_grad()
            output = model(X)
            # loss = criterion(output[-1, :, 0], y)
            loss = criterion(output[-1, :, 0].flatten(), y.flatten())
            loss.backward()
            optimizer.step()

    # 预测未来数据
    model.eval()
    predictions = []
    # last_sequence = torch.FloatTensor(scaled_data[-sequence_length:]).unsqueeze(1)
    last_sequence = torch.FloatTensor(scaled_data[-sequence_length:]).unsqueeze(0).unsqueeze(-1)

    for _ in range(92):  # 预测未来92天
        with torch.no_grad():
            pred = model(last_sequence.permute(1, 0, 2))[-1, :, 0].item()
            predictions.append(pred)
            last_sequence = torch.cat((last_sequence[1:], torch.tensor([[pred]])))

    # 将预测数据反归一化
    predicted_sales = np.array(predictions) * (scaler[1] - scaler[0]) + scaler[0]

    # 格式化预测数据并保存为JSON
    future_dates = pd.date_range(start='2023-07-01', periods=92, freq='D')
    forecast_data = [{"date": date.strftime('%Y-%m-%d'), "sales": sales} for date, sales in zip(future_dates, predicted_sales)]
    with open(output_json_path, 'w') as file:
        json.dump(forecast_data, file, indent=4)

    # 绘制图表
    plt.figure(figsize=(10, 6))
    plt.plot(df['ds'], df['y'], label='Original Sales', color='blue')
    plt.plot(future_dates, predicted_sales, label='Predicted Sales', color='orange')
    plt.xlabel('Date')
    plt.ylabel('Sales')
    plt.title('Sales Prediction with Transformer')
    plt.legend()
    plt.savefig(graph_output_path)
    plt.close()

# 单元测试主函数
if __name__ == "__main__":
    input_path = "../fujian/fujian2/groupByCategory/category_category1.json"
    output_path = "./unit_test/category_category1_pred.json"
    graph_path = "./unit_test/category_category1.png"
    predict_single_category(input_path, output_path, graph_path)
